# -*- coding: utf-8 -*-
# tensorflow packages
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.regularizers import l2
# other modules
from deepctr.inputs import build_input_features
from tensorflow.keras import Model, layers
from tensorflow.keras.layers import Concatenate
from tensorflow.keras.layers import Input, Embedding, Flatten, Dense

# my package
from .deepfm import DeepFM, DeepFMVariant
from deepctr.layers.sequence import SequencePoolingLayer
from .transformer import Encoder, create_padding_mask


class E2EModel:
    def __init__(self, max_len, max_time):
        self.max_len = max_len
        self.max_time = max_time

    def build_model(self, linear_feat_columns, feat_columns, sparse_feat_cols, dnn_hidden_units_1=(32, 32),
                    dnn_hidden_units_2=(32, 32), num_layers=1, num_heads=1, item_dim=32,
                    pos_trainable=False, l2_reg=0, dnn_dropout=0, dnn_use_bn=False,
                    use_fm=True, use_deep=True, dnn_activation='relu'
                    ):
        # be use of padding mask
        his_clk_items = Input(shape=(self.max_len,), dtype=tf.int32, name='his_clk_items')
        his_clk_time = Input(shape=(self.max_len,), dtype=tf.int32, name='his_clk_time')
        his_clk_times = Input(shape=(self.max_len,), dtype=tf.float32, name='his_clk_times')
        his_clk_items_with_profile = [build_input_features(feat_columns, 'his_item'+str(i)) for i in range(self.max_len)]
        inputs = [his_clk_items, his_clk_time, his_clk_times]
        # 输入为一个特征所有item
        for feat in sparse_feat_cols:
            for item in his_clk_items_with_profile:
                inputs.append(item[feat])

        deep_fm = DeepFM(linear_feat_columns, feat_columns,
                         dnn_hidden_units=dnn_hidden_units_1, output_dim=item_dim,
                         l2_reg=l2_reg, dnn_dropout=dnn_dropout, dnn_use_bn=dnn_use_bn,
                         dnn_activation=dnn_activation, use_fm=use_fm, use_deep=use_deep,
                         )
        user_embedding = [deep_fm(features) for features in his_clk_items_with_profile]
        user_embedding = Concatenate(axis=-2)(user_embedding)
        time_embedding = Embedding(self.max_time, user_embedding.shape[-1])
        his_clk_time_emb = time_embedding(his_clk_time)

        deep_fm_variant = DeepFMVariant(dnn_hidden_units=dnn_hidden_units_2, l2_reg=l2_reg,
                                        dnn_dropout=dnn_dropout, dnn_use_bn=dnn_use_bn,
                                        dnn_activation=dnn_activation, use_fm=use_fm, use_deep=use_deep,)
        time_item_list = []
        for i in range(self.max_len):
            time_item = Concatenate(axis=-2)([his_clk_time_emb[:, i:i+1, :], user_embedding[:, i:i+1, :]])
            time_item_emb = deep_fm_variant(time_item)
            time_item_list.append(time_item_emb)
        time_item_embs = Concatenate(axis=-2)(time_item_list)
        his_clk_times = tf.expand_dims(his_clk_times, axis=-1)
        feat_embs = tf.math.multiply(his_clk_times, time_item_embs)

        item_dim = feat_embs.shape[-1]
        # transformer model
        encoder = Encoder(num_layers=num_layers, d_model=item_dim, num_heads=num_heads,
                          dff=item_dim * 4, pos_trainable=pos_trainable, maximum_position_encoding=self.max_len)
        enc_padding_mask = create_padding_mask(his_clk_items)
        user_trm = encoder(feat_embs, enc_padding_mask)
        user_output = user_trm[:, -1:, :]
        user_output = Flatten()(user_output)

        age_pred = Dense(16, activation='relu')(user_output)
        gender_pred = Dense(16, activation='relu')(user_output)

        age_pred = Dense(10, activation='softmax', name='age')(age_pred)
        gender_pred = Dense(1, activation='sigmoid', name='gender')(gender_pred)

        model = Model(inputs=inputs, outputs=[age_pred, gender_pred])
        return model

    def build_model_2(self, linear_feat_columns, feat_columns, sparse_feat_cols, max_items, item_emb_matrix,
                      dnn_hidden_units_1=(32, 32), dnn_hidden_units_2=(32, 32), num_layers=1, num_heads=1,
                      item_dim=32, pos_trainable=False, l2_reg=0, dnn_dropout=0, dnn_use_bn=False,
                      use_fm=True, use_deep=True, dnn_activation='relu'
                      ):
        # be use of padding mask
        his_clk_items = Input(shape=(self.max_len,), dtype=tf.int32, name='his_clk_items')
        his_clk_time = Input(shape=(self.max_len,), dtype=tf.int32, name='his_clk_time')
        his_clk_times = Input(shape=(self.max_len,), dtype=tf.float32, name='his_clk_times')
        his_clk_items_with_profile = [build_input_features(feat_columns, 'his_item'+str(i)) for i in range(self.max_len)]
        inputs = [his_clk_items, his_clk_time, his_clk_times]
        # 输入为一个特征所有item
        for feat in sparse_feat_cols:
            for item in his_clk_items_with_profile:
                inputs.append(item[feat])

        deep_fm = DeepFM(linear_feat_columns, feat_columns,
                         dnn_hidden_units=dnn_hidden_units_1, output_dim=item_dim,
                         l2_reg=l2_reg, dnn_dropout=dnn_dropout, dnn_use_bn=dnn_use_bn,
                         dnn_activation=dnn_activation, use_fm=use_fm, use_deep=use_deep,
                         )
        user_embedding_1 = [deep_fm(features) for features in his_clk_items_with_profile]
        user_embedding_1 = Concatenate(axis=-2)(user_embedding_1)

        # item embedding
        item_embedding = Embedding(max_items, item_dim, weights=[item_emb_matrix], trainable=False)
        user_embedding_2 = item_embedding(his_clk_items)

        # user embedding 相加
        user_embedding = layers.Add()([user_embedding_1, user_embedding_2])

        time_embedding = Embedding(self.max_time, user_embedding.shape[-1])
        his_clk_time_emb = time_embedding(his_clk_time)

        deep_fm_variant = DeepFMVariant(dnn_hidden_units=dnn_hidden_units_2, l2_reg=l2_reg,
                                        dnn_dropout=dnn_dropout, dnn_use_bn=dnn_use_bn,
                                        dnn_activation=dnn_activation, use_fm=use_fm, use_deep=use_deep,)
        time_item_list = []
        for i in range(self.max_len):
            time_item = Concatenate(axis=-2)([his_clk_time_emb[:, i:i+1, :], user_embedding[:, i:i+1, :]])
            time_item_emb = deep_fm_variant(time_item)
            time_item_list.append(time_item_emb)
        time_item_embs = Concatenate(axis=-2)(time_item_list)
        his_clk_times = tf.expand_dims(his_clk_times, axis=-1)
        feat_embs = tf.math.multiply(his_clk_times, time_item_embs)

        item_dim = feat_embs.shape[-1]
        # transformer model
        encoder = Encoder(num_layers=num_layers, d_model=item_dim, num_heads=num_heads,
                          dff=item_dim * 4, pos_trainable=pos_trainable, maximum_position_encoding=self.max_len)
        enc_padding_mask = create_padding_mask(his_clk_items)
        user_trm = encoder(feat_embs, enc_padding_mask)
        age_pred = layers.Bidirectional(layers.LSTM(32))(user_trm)
        gender_pred = layers.Bidirectional(layers.LSTM(32))(user_trm)

        age_pred = Dense(16, activation='relu')(age_pred)
        gender_pred = Dense(16, activation='relu')(gender_pred)

        age_pred = Dense(10, activation='softmax', name='age')(age_pred)
        gender_pred = Dense(1, activation='sigmoid', name='gender')(gender_pred)

        model = Model(inputs=inputs, outputs=[age_pred, gender_pred])
        return model

    def build_model_3(self, linear_feat_columns, feat_columns, sparse_feat_cols, max_items, item_emb_matrix=None,
                      trainable=True, dnn_hidden_units_1=(32, 32), dnn_hidden_units_2=(32, 32), num_layers=1, num_heads=1,
                      item_dim=32, pos_trainable=False, l2_reg=0, dnn_dropout=0, dnn_use_bn=False, combiner='mean',
                      use_fm=True, use_deep=True, dnn_activation='relu'
                      ):
        # be use of padding mask
        his_clk_items = Input(shape=(self.max_len,), dtype=tf.int32, name='his_clk_items')
        his_clk_time = Input(shape=(self.max_len,), dtype=tf.int32, name='his_clk_time')
        his_clk_times = Input(shape=(self.max_len,), dtype=tf.float32, name='his_clk_times')
        user_seq_emb_1 = Input(shape=(self.max_len, item_dim), dtype=tf.float32, name='user_seq_emb_1')
        his_clk_items_with_profile = [build_input_features(feat_columns, 'his_item'+str(i)) for i in range(self.max_len)]
        inputs = [his_clk_items, his_clk_time, his_clk_times, user_seq_emb_1]
        # 输入为一个特征所有item
        for feat in sparse_feat_cols:
            for item in his_clk_items_with_profile:
                inputs.append(item[feat])

        deep_fm = DeepFM(linear_feat_columns, feat_columns,
                         dnn_hidden_units=dnn_hidden_units_1, output_dim=item_dim,
                         l2_reg=l2_reg, dnn_dropout=dnn_dropout, dnn_use_bn=dnn_use_bn,
                         dnn_activation=dnn_activation, use_fm=use_fm, use_deep=use_deep,
                         )
        user_seq_emb_2 = [deep_fm(features) for features in his_clk_items_with_profile]
        user_seq_emb_2 = Concatenate(axis=-2)(user_seq_emb_2)

        # item embedding
        # if item_emb_matrix is not None:
        #     item_embedding = Embedding(max_items, item_dim, weights=[item_emb_matrix], trainable=False)
        # else:
        #     item_embedding = Embedding(max_items, item_dim, embeddings_initializer=keras.initializers.Zeros(),
        #                                embeddings_regularizer=l2(l2_reg), trainable=trainable)
        #
        # user_embedding_2 = item_embedding(his_clk_items)

        # user embedding 相加
        user_embedding = layers.Add()([user_seq_emb_1, user_seq_emb_2])

        # time_embedding = Embedding(self.max_time, user_embedding.shape[-1],
        #                            embeddings_regularizer=l2(l2_reg), trainable=trainable)
        # his_clk_time_emb = time_embedding(his_clk_time)
        #
        # deep_fm_variant = DeepFMVariant(dnn_hidden_units=dnn_hidden_units_2, l2_reg=l2_reg,
        #                                 dnn_dropout=dnn_dropout, dnn_use_bn=dnn_use_bn,
        #                                 dnn_activation=dnn_activation, use_fm=use_fm, use_deep=use_deep,)
        # time_item_list = []
        # for i in range(self.max_len):
        #     time_item = Concatenate(axis=-2)([his_clk_time_emb[:, i:i+1, :], user_embedding[:, i:i+1, :]])
        #     time_item_emb = deep_fm_variant(time_item)
        #     time_item_list.append(time_item_emb)
        # time_item_embs = Concatenate(axis=-2)(time_item_list)
        his_clk_times = tf.expand_dims(his_clk_times, axis=-1)
        feat_embs = tf.math.multiply(his_clk_times, user_embedding)

        item_dim = feat_embs.shape[-1]
        # transformer model
        encoder = Encoder(num_layers=num_layers, d_model=item_dim, num_heads=num_heads,
                          dff=item_dim * 4, pos_trainable=pos_trainable, maximum_position_encoding=self.max_len)
        enc_padding_mask = create_padding_mask(his_clk_items)
        user_trm = encoder(feat_embs, enc_padding_mask)

        masking_layer = layers.Masking()
        masked_user_trm = masking_layer(user_trm)
        pool_vec = SequencePoolingLayer(combiner, supports_masking=True)(masked_user_trm)
        pool_vec = Flatten()(pool_vec)

        age_pred = Dense(32, activation='relu')(pool_vec)
        age_pred = Dense(16, activation='relu')(age_pred)
        age_pred = Dense(10, activation='softmax', name='age')(age_pred)

        gender_pred = Dense(32, activation='relu')(pool_vec)
        gender_pred = Dense(16, activation='relu')(gender_pred)
        gender_pred = Dense(8, activation='relu')(gender_pred)
        gender_pred = Dense(4, activation='relu')(gender_pred)
        gender_pred = Dense(2, activation='relu')(gender_pred)
        gender_pred = Dense(1, activation='sigmoid', name='gender')(gender_pred)

        model = Model(inputs=inputs, outputs=[age_pred, gender_pred])
        return model



